refactor: pass SubqueryContext explicitly through planner traits#23649
refactor: pass SubqueryContext explicitly through planner traits#23649timsaucer wants to merge 5 commits into
Conversation
AI Disclosure: This code was written in part by an AI agent.
github-actions
Bot
added
documentation
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Replace broad uses of expression lowering with explicit language about creating physical expressions, and correct the migration guide PR link. AI Disclosure: This code was written in part by an AI agent.
timsaucer
left a comment
timsaucer
left a comment
There was a problem hiding this comment.
I've tried to mark all of the places in the code that I think are significant, which are few. The rest is just adding in an extra parameter to function calls and plumbing that through as necessary.
| let mut owned = session_state.clone(); | ||
| owned.execution_props_mut().subquery_indexes = index_map; | ||
| owned.execution_props_mut().subquery_results = results.clone(); | ||
| let session_state = Cow::Owned(owned); |
There was a problem hiding this comment.
This is the work around we are removing.
| #[tokio::test] | ||
| async fn scalar_subquery_in_extension_expr_plans() -> Result<()> { | ||
| let subquery = LogicalPlanBuilder::empty(true) | ||
| .project(vec![lit(42_i32)])? | ||
| .build()?; | ||
| let logical_plan = LogicalPlan::Extension(Extension { | ||
| node: Arc::new(NoOpExtensionNode { | ||
| expressions: vec![scalar_subquery(Arc::new(subquery))], | ||
| ..Default::default() | ||
| }), | ||
| }); | ||
| let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( | ||
| ExpressionExtensionPlanner, | ||
| )]); | ||
|
|
||
| let plan = planner | ||
| .create_physical_plan(&logical_plan, &make_session_state()) | ||
| .await?; | ||
|
|
||
| assert_contains!(format!("{plan:?}"), "ScalarSubqueryExec"); | ||
| Ok(()) | ||
| } | ||
|
|
There was a problem hiding this comment.
This test is necessary and demonstrates extension nodes that need the subquery context piped through properly.
| subquery_indexes: HashMap::new(), | ||
| subquery_results: ScalarSubqueryResults::default(), |
There was a problem hiding this comment.
This is the core of the change to this PR. We move these values out of the ExecutionProps and into their own struct, SubqueryContext
|
@gabotechs I think this is the cleanest, non-hacky way to actually plumb through the subquery context in the logical to physical conversion. It's also the largest change, and API breaking, but the update to downstream users is pretty straight forward I think. Also tagging @neilconway for the bit about the subquery context. |
There was a problem hiding this comment.
👍 No surprises, this is mostly what I already reviewed in #22340. Looks good.
It's a shame that we need this change to be this big, but after trying several other alternatives (#23418, #23622) it seems appropriate and justified.
It's also nice to avoid a clone of the SessionState.
From what we've tried, this is the only non-brittle solution, and the one that looks more future proof, so +1 to it.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23649 +/- ##
========================================
Coverage 80.67% 80.67%
========================================
Files 1086 1086
Lines 366800 367068 +268
Branches 366800 367068 +268
========================================
+ Hits 295912 296150 +238
- Misses 53260 53294 +34
+ Partials 17628 17624 -4 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you for this PR @timsaucer and @gabotechs
I agree this is the most expedient way to proceed and I think it is ok.
I spent a while reviewing the code and trying to figure out why we need something different than SessionContext, ExecutionProps and TaskContext which are supposed to track the per Session, per plan, and per-execution state
I think the core difference (that this PR is filling) is that we have no structure that is scoped to some subtree of the plan (rather than the entire plan), and I think SubqueryContext fills this gap -- we have something similar in the sql planner to handle subqueries here:
datafusion/datafusion/sql/src/planner.rs
Line 289 in ca02890
However, I would like to recommend:
- Rename
SubqueryContextto something more general likePhysicalPlanningContext - Move it to its own module (like
planning_context.rsorphysical_planning_context.rs)
As a follow on perhaps we can move the lambda handing in there as well
datafusion/datafusion/physical-expr/src/planner.rs
Lines 612 to 614 in 44c8305
|
|
||
| ### Scalar-subquery state moved from `ExecutionProps` to an explicit `SubqueryContext` parameter | ||
|
|
||
| The `subquery_indexes` and `subquery_results` public fields on |
| let physical_expr = datafusion::physical_expr::create_physical_expr( | ||
| &expr, | ||
| &df_schema, | ||
| &props, |
There was a problem hiding this comment.
One question I have is can we possible put the subquery context as a field on ExecutionProps (as it seems specific to each execution -- which is supposed to be the scope of the ExecutionProps)
| use datafusion::execution::context::QueryPlanner; | ||
| use datafusion::execution::session_state::CacheFactory; | ||
| use datafusion::execution::{SessionState, SessionStateBuilder}; | ||
| use datafusion::logical_expr::execution_props::SubqueryContext; |
There was a problem hiding this comment.
I found it confusing that SubqueryContext is logically part of ExecutionProps (which is supposed to be per-execution state) but SubqueryContext has
| /// [`Subquery`]: crate::logical_plan::Subquery | ||
| #[derive(Clone, Debug, Default)] | ||
| pub struct SubqueryContext { | ||
| indexes: HashMap<crate::logical_plan::Subquery, SubqueryIndex>, |
There was a problem hiding this comment.
I find this somewhat confusing that the subquery context has both planning time info (indexes) and something that is modified during execution time (results)
There was a problem hiding this comment.
Also it is not clear to me why lambda expressions don't have the same issue 🤔
There was a problem hiding this comment.
Ah, they actually do:
datafusion/datafusion/physical-expr/src/planner.rs
Lines 612 to 614 in 44c8305
| } | ||
| } | ||
|
|
||
| /// Per-plan context used by the physical planner when creating physical |
There was a problem hiding this comment.
Technically speaking I don't think this is "per plan" -- instead it is more like "per plan subtree"
4 participants
Which issue does this PR close?
Rationale for this change
When we turn logical plans into physical plans, we have a work around to get the context for scalar subqueries. Specifically we clone the session and mutate the
execution_propswhere the subquery context currently sits. This is even called out in the code with comments:The real reason for this is that we want to expose the
QueryPlannervia FFI and currently it depends upon downcasts fromSessiontoSessionState. That does not work for the FFI crate where we have a differentSessionimplementation. Since we cannot do the downcast, we cannot use the work around in the current code.This PR is designed to plumb through the required context properly rather than rely on manipulating the execution properties.
What changes are included in this PR?
The major change here is to remove
subquery_indexesandsubquery_resultsfromExecutionPropsand put them into their own structSubqueryContext. Everything else is just refactoring and plumbing to align around this split of the data.Are these changes tested?
Are there any user-facing changes?
Yes, this adds a single parameter into
create_physical_expr, the subquery context. The migration guide demonstrates how to add a default context in for users, if necessary.